#!/usr/bin/env python3
import os

from data_utils import *
from publish_utils import *
from kitti_util import *
from head_function import *
from tracklet_parser import parse_xml
from collections import deque

DATA_PATH = '/home/jony/data/kitti/RawData/2011_09_26/2011_09_26_drive_0061_sync/'
EGOCAR = np.array([[2.15, 0.9, -1.73], [2.15, -0.9, -1.73], [-1.95, -0.9, -1.73], [-1.95, 0.9, -1.73],
                    [2.15, 0.9, -0.23], [2.15, -0.9, -0.23], [-1.95, -0.9, -0.23], [-1.95, 0.9, -0.23]])
# 本车在点云坐标系下的3D检测框底面0123的交点
EGOCAR_POINT = [0.1, 0, -1.73]
def compute_distance_from_velocity(vf, vl, vu, time_interval):
    """
    Compute the displacement vector from velocities.
    """
    vf, vl, vu = vf.iloc[0], vl.iloc[0], vu.iloc[0]
    displacement = np.array([vf * time_interval, vl * time_interval, vu * time_interval])
    return displacement

class Object():
    def __init__(self, center):
        """
        初始化物体对象，用于存储物体的轨迹位置信息，设置轨迹位置队列最大长度为20
        """
        self.locations = deque(maxlen=20)
        self.locations.appendleft(center)

    def update(self, center, displacement, yaw_change, pitch_change, roll_change):
        """
        根据给定的中心位置、位移、偏航角、俯仰角和横滚角变化更新物体的所有轨迹位置信息，
        若传入新的中心位置不为None，则添加到轨迹位置队列头部
        """
        for i in range(len(self.locations)):
            x0, y0, z0 = self.locations[i]
            # 利用旋转矩阵进行姿态变化更新，先根据偏航角变化更新
            x1 = x0 * np.cos(yaw_change) + y0 * np.sin(yaw_change)
            y1 = -x0 * np.sin(yaw_change) + y0 * np.cos(yaw_change)
            # 再根据俯仰角变化更新
            x2 = x1 * np.cos(pitch_change) - z0 * np.sin(pitch_change)
            z1 = x1 * np.sin(pitch_change) + z0 * np.cos(pitch_change)
            # 最后根据横滚角变化更新
            y2 = y1 * np.cos(roll_change) + z1 * np.sin(roll_change)
            z2 = -y1 * np.sin(roll_change) + z1 * np.cos(roll_change)
            # 减去位移量得到更新后的位置
            x3 = x2 - displacement[0]
            y3 = y2 - displacement[1]
            z3 = z2 - displacement[2]
            self.locations[i] = np.array([x3, y3, z3])
        
        if center is not None:
            self.locations.appendleft(center)

    def reset(self):
        """
        重置物体的轨迹位置信息，清空轨迹位置并保留最大长度限制
        """
        self.locations = deque(maxlen=20)


# class Object():
#     def __init__(self, center):
#         self.locations = deque(maxlen=20)
#         self.locations.appendleft(center)

#     def update(self, center, displacement, yaw_change):
#         for i in range(len(self.locations)):
#             x0, y0 = self.locations[i]
#             x1 = x0 * np.cos(yaw_change) + y0 * np.sin(yaw_change) - displacement[0]
#             y1 = -x0 * np.sin(yaw_change) + y0 * np.cos(yaw_change)- displacement[1]
#             self.locations[i] = np.array([x1, y1])
        
#         if center is not None:
#             self.locations.appendleft(center)

#     def reset(self):
#         self.locations = deque(maxlen=20)  # 清空轨迹位置并保留最大长度限制


if __name__ == '__main__':
    frame = 0
    rospy.init_node('kitti_node', anonymous=True)
    cam_pub = rospy.Publisher('kitti_cam', Image, queue_size=10)
    pcl_pub = rospy.Publisher('kitti_pcl', PointCloud2, queue_size=10)
    ego_pub = rospy.Publisher('kitti_ego_car', MarkerArray, queue_size=10)
    imu_pub = rospy.Publisher('kitti_imu', Imu, queue_size=10)
    gps_pub = rospy.Publisher('kitti_gps', NavSatFix, queue_size=10)
    box3d_pub = rospy.Publisher('kitti_3d', MarkerArray, queue_size=10)
    track_pub = rospy.Publisher('kitti_track', MarkerArray, queue_size=10)
    distance_pub = rospy.Publisher('kitti_distance', MarkerArray, queue_size=10)
    bridge = CvBridge()

    rate = rospy.Rate(10)

    tracklets = parse_xml('/home/jony/data/kitti/tracklets_0061/2011_09_26/2011_09_26_drive_0061_sync/tracklet_labels.xml')
    calib = Calibration('/home/jony/data/kitti/RawData/2011_09_26/', from_video=True)
    df_tracklets = get_frame_fields(tracklets)

    tracker = {}  # track_id : Object
    prev_imu_data = None

    while not rospy.is_shutdown():
        time_interval = 0.1
        df_tracklets_frame = df_tracklets[df_tracklets.frame == frame]
        types = np.array(df_tracklets_frame['type'])
        boxes_3d = np.array(df_tracklets_frame[['h', 'w', 'l', 'tx', 'ty', 'tz', 'rz']])
        track_ids = np.array(df_tracklets_frame['tracklet_idx'])
        boxes_2d = []
        corners_3d_velos = []
        centers = {}  # track_id: center
        minPQDs = [] # 与自身车辆的最小距离
        for track_id, box_3d in zip(track_ids, boxes_3d):
            h, w, l = box_3d[0], box_3d[1], box_3d[2]
            tx, ty, tz = box_3d[3], box_3d[4], box_3d[5]
            rz = wrap_to_pi(box_3d[6])
            yaw = rz_to_ry(rz)
            # 将世界坐标系下的 tx, ty, tz 转换为相机坐标系下的 pos_x, pos_y, pos_z
            pos_x, pos_y, pos_z = convert_world_to_camera_coords(tx, ty, tz, calib)
            # 使用转换后的坐标计算3D框
            corners_2d, corners_3d_cam2 = compute_3d_box_cam2(h, w, l, pos_x, pos_y, pos_z, yaw, calib)
            corners_3d_velo = calib.project_rect_to_velo(corners_3d_cam2.T)
            box_2d = get_2d_boxes(corners_2d)
            minPQD = min_distance_cuboids(EGOCAR, corners_3d_velo)
            boxes_2d.append(box_2d)
            corners_3d_velos.append(corners_3d_velo)
            minPQDs.append(minPQD)
            # 垂直方向对检测框(x,y)取平均
        #     centers[track_id] = np.mean(corners_3d_velo, axis = 0)[:2]
        # centers[-1] = np.array([0, 0])  # 添加自身车辆的行车轨迹
            centers[track_id] = np.mean(corners_3d_velo, axis=0)[:3]
        centers[-1] = np.array([0, 0, 0])  # 添加自身车辆的行车轨迹

        image = read_camera(os.path.join(DATA_PATH, 'image_02/data/%010d.png' % frame))
        point_cloud = read_point_cloud(os.path.join(DATA_PATH, 'velodyne_points/data/%010d.bin' % frame))
        imu_data = read_imu(os.path.join(DATA_PATH, 'oxts/data/%010d.txt' % frame))

        if prev_imu_data is None:
            for track_id in centers:
                tracker[track_id] = Object(centers[track_id])
        else:
            displacement = compute_distance_from_velocity(imu_data.vf, imu_data.vl, imu_data.vu, time_interval)
            yaw_change = imu_data.yaw.iloc[0] - prev_imu_data.yaw.iloc[0]
            pitch_change = imu_data.pitch.iloc[0] - prev_imu_data.pitch.iloc[0]
            roll_change = imu_data.roll.iloc[0] - prev_imu_data.roll.iloc[0]
            for track_id in centers:
                # 若当前帧中被追踪目标，上一帧就存在，则对应更新其中心位置，dispalcement, 以及各角度变化
                if track_id in tracker:
                    tracker[track_id].update(centers[track_id], displacement, yaw_change, pitch_change, roll_change)
                else:
                    tracker[track_id] = Object(centers[track_id])  # 反之，第一次被追踪，则记录该目标（物体）
            for track_id in tracker:  # 上一帧存在，本帧中被遮挡（不存在）
                if track_id not in centers:
                    tracker[track_id].update(None, displacement, yaw_change, pitch_change, roll_change)

        prev_imu_data = imu_data

        publish_camera(cam_pub, bridge, image, boxes_2d, types)
        publish_point_cloud(pcl_pub, point_cloud)
        publish_ego_car(ego_pub)
        publish_imu(imu_pub, imu_data)
        publish_gps(gps_pub, imu_data)
        publish_3dbox(box3d_pub, corners_3d_velos, types)
        publish_track(track_pub, tracker, centers)
        publish_distance(distance_pub, minPQDs)
        rospy.loginfo("Published: frame %d" % frame)
        rate.sleep()
        frame += 1 # 遍历每张图片
        if frame == 703:
            frame = 0
            for track_id in tracker:
                tracker[track_id].reset()
